import numpy as np
import cv2
from scipy.signal import fftconvolve

def generate_donut_kernel(inner_radius, outer_radius):
    """
    生成一个环形核 (donut kernel)，用于增强圆心检测。
    """
    size = 2 * outer_radius + 1
    y, x = np.ogrid[-outer_radius:outer_radius+1, -outer_radius:outer_radius+1]
    dist2 = x**2 + y**2
    
    kernel = np.zeros((size, size), dtype=np.float32)
    kernel[(dist2 >= inner_radius**2) & (dist2 <= outer_radius**2)] = 1.0
    return kernel


def locate_pupil_center(img, bandwidth):
    """
    使用2D核密度估计(FFT变体)和环形核定位瞳孔中心
    """
    # 创建圆柱形核
    kernel = generate_donut_kernel(inner_radius=bandwidth//2, outer_radius=bandwidth)
    
    # 使用 FFT 加速卷积
    density_map = fftconvolve(img, kernel, mode='same')
    
    # 找到峰值位置作为圆心
    return np.unravel_index(np.argmax(density_map), density_map.shape)

def create_pupil_mask(img, center, bandwidth):
    mask = np.zeros_like(img, dtype=np.uint8)
    cv2.circle(mask, (center[0], center[1]), bandwidth, 255, -1)
    return mask.astype(np.bool)